/*******************************************************************************
**                      Includes                                              **
*******************************************************************************/
#include <string.h>
#include <stdlib.h>
#include <stdio.h>

#include "sm3.h"
/*******************************************************************************
**                      Private Macro Definitions                             **
*******************************************************************************/
/*
 * 32-bit integer manipulation macros (big endian)
 */
#ifndef GET_ULONG_BE
#define GET_ULONG_BE(n,b,i)                      \
{                                                \
	(n) = ( (uint32) (b)[(i)    ] << 24 )   \
		| ( (uint32) (b)[(i) + 1] << 16 )   \
		| ( (uint32) (b)[(i) + 2] <<  8 )   \
		| ( (uint32) (b)[(i) + 3]       );  \
}
#endif

#ifndef PUT_ULONG_BE
#define PUT_ULONG_BE(n,b,i)                   \
{                                             \
	(b)[(i)    ] = (uint8) ( (n) >> 24 );    \
	(b)[(i) + 1] = (uint8) ( (n) >> 16 );    \
	(b)[(i) + 2] = (uint8) ( (n) >>  8 );    \
	(b)[(i) + 3] = (uint8) ( (n)       );    \
}
#endif

#define FF0(x,y,z) ( (x) ^ (y) ^ (z))
#define FF1(x,y,z) (((x) & (y)) | ( (x) & (z)) | ( (y) & (z)))

#define GG0(x,y,z) ( (x) ^ (y) ^ (z))
#define GG1(x,y,z) (((x) & (y)) | ( (~(x)) & (z)) )


#define SHL(x,n) (((x) & 0xFFFFFFFF) << (n % 32))
// #define ROTL(x,n) (SHL((x),n) | ((x) >> (32 - n)))
#define ROTL(x,n) (SHL((x),n) | ((x) >> (32 - (n % 32))))

#define P0(x) ((x) ^  ROTL((x),9) ^ ROTL((x),17))
#define P1(x) ((x) ^  ROTL((x),15) ^ ROTL((x),23))

static uint8 sm3_padding[64] =
{
	0x80, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
	0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
	0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
	0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
};

/*******************************************************************************
**                      Private Function Declarations                         **
*******************************************************************************/

/************************************************************************************************************
 *  Cry_SM3_BiToW()
 ***********************************************************************************************************/
/*!
 * \brief       Expands a 64-byte data block into 68 32-bit words following the SM3 specification.
 * \details     - Validates input parameters (non-NULL pointers, sufficient buffer sizes).
 *              - Reads the first 16 words (W[0..15]) from the 64-byte input buffer \p Bi in big-endian form.
 *              - Generates the remaining 52 words (W[16..67]) using the SM3 expansion formula:
 *                \code
 *                tmp = W[i - 16] ^ W[i - 9] ^ ROTL(W[i - 3], 15);
 *                W[i] = P1(tmp) ^ ROTL(W[i - 13], 7) ^ W[i - 6];
 *                \endcode
 *
 * \param[in]   Bi      Pointer to the input buffer containing 64 bytes of SM3 block data.
 * \param[in]   BiLen   Length of the input buffer in bytes (must be at least 64).
 * \param[out]  W       Pointer to the output array of 68 32-bit words.
 * \param[in]   WLen    The number of 32-bit words available in the output array (must be at least 68).
 *
 * \return      E_OK        If the expansion was successful.
 * \return      E_NOT_OK    If the parameters are invalid (e.g., NULL pointers or insufficient buffer sizes).
 *
 * \pre         The caller must ensure \p Bi is at least 64 bytes and \p W can hold at least 68 words.
 * \context     Any
 * \reentrant   No
 * \synchronous Yes
 ***********************************************************************************************************/
static Std_ReturnType Cry_SM3_BiToW(uint8_t *Bi, size_t BiLen, uint32_t *W, size_t WLen)
{
    uint8 i;
    uint32 tmp;

    /* 1. Parameter validation */
    if ((Bi == NULL) || (W == NULL)) 
    {
        return E_NOT_OK;
    }

    if (BiLen < 64 || WLen < 68) 
    {         
        return E_NOT_OK;     
    }

    /* 2. Load the first 16 words (W[0..15]) from the input buffer */
    for (i = 0; i < 16; i++) 
    {
        GET_ULONG_BE(W[i], Bi, i * 4);
    }

    /* 3. Compute W[16..67] using the SM3 message expansion formula */
    for (i = 16; i < 68; i++) 
    {
        tmp = W[i - 16]    ^ W[i - 9] ^ ROTL(W[i - 3], 15);
        W[i] = P1(tmp) ^ (ROTL(W[i - 13], 7)) ^ W[i - 6];
    }

    return E_OK;
}

/************************************************************************************************************
 *  Cry_SM3_WToW1()
 ***********************************************************************************************************/
/*!
 * \brief         Extend the SM3 expanded message W[] to produce W1[].
 * \details       This function computes an additional 64 32-bit words (\a W1) based on the 
 *                expanded message array (\a W) using the formula:
 *                \f$ W1[i] = W[i] \oplus W[i+4] \f$
 *                where \( i \in [0, 63] \). The computed \a W1[] is then used during the 
 *                64 compression rounds of the SM3 algorithm.
 *
 * \param[in]     W          Array of 32-bit expanded words (\a W[0..67]), size must be at least 68.
 * \param[out]    W1         Output array to store 64 extended words (\a W1[0..63]), size must be at least 64.
 * \param[in]     WLen       Length of the input array \a W, must be at least 68.
 * \param[in]     W1Len      Length of the output array \a W1, must be at least 64.
 *
 * \return        Std_ReturnType
 *                - \c E_OK:     Operation completed successfully.
 *                - \c E_NOT_OK: Invalid input parameters (e.g., NULL pointers or insufficient array size).
 *
 * \pre           - \a W and \a W1 must be valid non-null pointers.
 *                - \a W must have at least 68 elements.
 *                - \a W1 must have at least 64 elements.
 *
 * \context       ANY
 * \reentrant     FALSE
 * \synchronous   TRUE
 * \config        -
 **********************************************************************************************************/
static Std_ReturnType Cry_SM3_WToW1(uint32 W[], uint32 W1[], size_t WLen, size_t W1Len)
{
    int i;

    // 1. Parameter validation
    if (W == NULL || W1 == NULL || WLen < 68 || W1Len < 64) 
    {
        return E_NOT_OK; // Return error if inputs are invalid
    }

    // 2. Compute W1[] based on the formula W1[i] = W[i] ^ W[i+4]
    for (i = 0; i < 64; i++)
    {
        W1[i] = W[i] ^ W[i + 4];
    }

    // 3. Indicate success
    return E_OK; 
}

/************************************************************************************************************
 *  Cry_SM3_Compress()
 ***********************************************************************************************************/
/*!
 * \brief         Perform the SM3 compression function (64-round core).
 * \details       This function implements the 64-round core compression logic of the SM3 algorithm. 
 *                It combines the expanded message arrays \a W[] and \a W1[] with the current hash state 
 *                to update the internal hash context (\a ctx->state).
 *
 * \param[in]     W          Input array of 68 32-bit words derived from the expanded message.
 *                           \a W must contain at least 68 elements.
 * \param[in]     W1         Input array of 64 32-bit words derived from \a W[] using the formula:
 *                           \f$ W1[i] = W[i] \oplus W[i+4], \, i \in [0, 63] \f$.
 *                           \a W1 must contain at least 64 elements.
 * \param[in,out] ctx        Pointer to the SM3 context structure, which includes:
 *                           - \c ctx->state: The current 256-bit hash state, updated after compression.
 *
 * \return        Std_ReturnType
 *                - \c E_OK:     Operation completed successfully.
 *                - \c E_NOT_OK: Invalid input parameters (e.g., NULL pointers).
 *
 * \pre           - \a W and \a W1 must be valid, non-null pointers.
 *                - \a ctx must be a valid, non-null pointer with properly initialized \c ctx->state.
 *                - \a W must have at least 68 elements, and \a W1 must have at least 64 elements.
 *
 * \context       ANY
 * \reentrant     FALSE
 * \synchronous   TRUE
 * \config        -
 **********************************************************************************************************/

static Std_ReturnType Cry_SM3_Compress(const uint32 * W, const uint32 * W1, sm3_ctx * ctx)
{
    uint32 SS1 = 0, SS2 = 0, TT1 = 0, TT2 = 0;
    // Temporary state variables
    uint32 A, B, C, D, E, F, G, H;
    // Round constant
    uint32 Tj = 0; 
    // Loop counter
    uint8 j;

    // 1. Parameter validation: Ensure non-null pointers
    if (W == NULL || W1 == NULL || ctx == NULL) 
    {
        return E_NOT_OK; // Fail silently or log an error if needed
    }

    // 2. Load the current state into local variables for faster access
    A = ctx->state[0];
    B = ctx->state[1];
    C = ctx->state[2];
    D = ctx->state[3];
    E = ctx->state[4];
    F = ctx->state[5];
    G = ctx->state[6];
    H = ctx->state[7];

    // 3. Perform 64 compression rounds
    for (j = 0; j < 64; j++)
    {
        // Compute the round constant Tj (use precomputed constants if possible)
        Tj = (j < 16) ? SM3_T_0 : SM3_T_1;

        // Compute SS1 and SS2
        SS1 = ROTL((ROTL(A, 12) + E + ROTL(Tj, j)), 7);
        SS2 = SS1 ^ ROTL(A, 12);

        // Compute TT1 and TT2 using the appropriate function based on the round
        if (j < 16)
        {
            TT1 = FF0(A, B, C) + D + SS2 + W1[j];
            TT2 = GG0(E, F, G) + H + SS1 + W[j];
        }
        else
        {
            TT1 = FF1(A, B, C) + D + SS2 + W1[j];
            TT2 = GG1(E, F, G) + H + SS1 + W[j];
        }

        // Update state variables for the next round
        D = C;
        C = ROTL(B, 9);
        B = A;
        A = TT1;
        H = G;
        G = ROTL(F, 19);
        F = E;
        E = P0(TT2);
    }

    // 5. Update the context state by XORing with the temporary variables
    ctx->state[0] ^= A;
    ctx->state[1] ^= B;
    ctx->state[2] ^= C;
    ctx->state[3] ^= D;
    ctx->state[4] ^= E;
    ctx->state[5] ^= F;
    ctx->state[6] ^= G;
    ctx->state[7] ^= H;

    // 6. Indicate success
    return E_OK; 
}


/************************************************************************************************************
 *  Cry_SM3_Process()
 ***********************************************************************************************************/
/*!
 * \brief         SM3 compression function (core compression round).
 * \details       This function implements the core logic of the SM3 hash algorithm. 
 *                It takes the 512-bit message block stored in \a ctx->buf and performs:
 *                1) Expansion into 68 32-bit words (\a W),
 *                2) Extension to derive another 64 words (\a W1),
 *                3) 64-round combination/update of the internal state in \c ctx->state.
 *
 * \param[in,out] ctx        Pointer to the SM3 context structure, holding:
 *                           - \c ctx->buf: a 64-byte buffer (512 bits) to be compressed,
 *                           - \c ctx->state: the current 256-bit SM3 state.
 *
 * \return        void
 *
 * \pre           - \a ctx must be valid and non-null.
 *                - \a ctx->buf must contain the 64 bytes of data to be compressed.
 * \context       ANY
 * \reentrant     FALSE
 * \synchronous   TRUE
 * \config        -
 **********************************************************************************************************/
static Std_ReturnType Cry_SM3_Process(sm3_ctx *ctx)
{
    /* 1. Local arrays for SM3 expansion/extension:
     * W  - holds 68 words derived from ctx->buf (512 bits).
     * W1 - holds 64 words derived from W[]. */
    uint32 W[68];
    uint32 W1[64];

    /* 2. Parameter check: ensure the input pointer is not NULL */
    if (ctx == NULL) 
    {
        return E_NOT_OK;
    }

    /* 3. Check 4-byte alignment for both the context buffer and input data */
    if (((uintptr_t)(ctx->buf) % 4) != 0) 
    {
        return E_NOT_OK;
    }

    /* 4. Expand the 64-byte block (ctx->buf) into 68 words (W).
     *    This is based on the SM3 specification's message expansion. */
    if(Cry_SM3_BiToW(ctx->buf, 64, W, 68) != E_OK)
    {
        return E_NOT_OK;
    }

    /* 5. Extend W[] to produce an additional 64 words (W1[]).  W1[i] = W[i] ^ W[i+4]. */
    if(Cry_SM3_WToW1(W, W1, 68, 64) != E_OK)
    {
        return E_NOT_OK;
    }

    /* 6. Compress the expanded data with the current state. Performs 64 rounds of SM3's main loop, updating ctx->state. */
    if(Cry_SM3_Compress(W, W1, ctx) != E_OK)
    {
        return E_NOT_OK;
    }

    return E_OK;
}


/************************************************************************************************************
 *  Cry_SM3_Init()
 ***********************************************************************************************************/
/*!
 * \brief         Initializes the SM3 context.
 * \details       This function initializes the SM3 context structure, setting the internal state to the 
 *                SM3 initial vector (IV) as defined by the SM3 standard. It also clears the message buffer
 *                and resets the buffer length and compressed data length.
 *
 * \param[in,out] ctx         Pointer to the SM3 context structure to be initialized.
 *
 * \return        E_OK        The initialization completed successfully.
 * \return        E_NOT_OK    The initialization failed (e.g., NULL pointer input).
 *
 * \pre           The calling environment must provide a valid `ctx` pointer.
 * \context       ANY
 * \reentrant     FALSE
 * \synchronous   TRUE
 * \config        -
 **********************************************************************************************************/
Std_ReturnType Cry_SM3_Init(sm3_ctx *ctx)
{
    /* 1. Parameter check: ensure the input pointer is not NULL */
    if (ctx == NULL) 
    {
        return E_NOT_OK; /* Return an error if ctx is NULL */
    }

    /* 2. Initialize the internal state of SM3 with the predefined IV */
    ctx->state[0] = SM3_IV_A;
    ctx->state[1] = SM3_IV_B;
    ctx->state[2] = SM3_IV_C;
    ctx->state[3] = SM3_IV_D;
    ctx->state[4] = SM3_IV_E;
    ctx->state[5] = SM3_IV_F;
    ctx->state[6] = SM3_IV_G;
    ctx->state[7] = SM3_IV_H;

    /* 3. Initialize the buffer length and compressed length to zero */
    ctx->cur_buf_len = 0;
    ctx->compressed_len = 0;

    /* 4. Return success */
    return E_OK;
}


/************************************************************************************************************
 *  Cry_SM3_Update()
 ***********************************************************************************************************/
/*!
 * \brief         Updates the SM3 context with new input data.
 * \details       This function continuously appends input data to the internal buffer \a ctx->buf. Whenever
 *                the buffer size reaches 64 bytes, it immediately calls the SM3 compression function 
 *                (\c Cry_SM3_Process) to process the data. The total amount of compressed data is tracked 
 *                in bits by \a ctx->compressed_len.
 *
 * \param[in,out] ctx        Pointer to the SM3 context structure.
 *                           - On entry, \a ctx->cur_buf_len indicates how many bytes are currently stored 
 *                             in the buffer.
 *                           - On exit, \a ctx->cur_buf_len is updated to reflect any remaining unprocessed bytes.
 *
 * \param[in]     input      Pointer to the input data to be hashed.
 *
 * \param[in]     iLen       The length of the input data (in bytes).
 *
 * \return        E_OK       Operation completed successfully.
 * \return        E_NOT_OK   One or more parameters are invalid (e.g., NULL pointer or misaligned addresses).
 *
 * \pre           - \a ctx and \a input must be valid, non-NULL pointers.
 *                - \a ctx->buf and \a input must be properly 4-byte aligned.
 * \context       ANY
 * \reentrant     FALSE
 * \synchronous   TRUE
 * \config        -
 **********************************************************************************************************/
Std_ReturnType Cry_SM3_Update(sm3_ctx *ctx, const uint8 *input, uint32 iLen)
{

    /* 1. Check for NULL pointers */
    if (ctx == NULL || input == NULL) 
    {
        return E_NOT_OK;
    }

    /* 2. Check 4-byte alignment for both the context buffer and input data */
    if (((uintptr_t)(ctx->buf) % 4) != 0 || ((uintptr_t)(input) % 4) != 0) 
    {
        return E_NOT_OK;
    }

    /* 3. Process each byte of the input data */
	while (iLen--)
    {
        ctx->buf[ctx->cur_buf_len] = *input++;
        ctx->cur_buf_len++;

        /* If the buffer is full (64 bytes), call the compression function */
        if (ctx->cur_buf_len == 64)
        {
            if (Cry_SM3_Process(ctx) != E_OK) 
            {
                return E_NOT_OK;  // failed to Compress
            }

            /* Update total compressed length in bits (64 bytes * 8 = 512 bits) */
            ctx->compressed_len += 512;

            /* Reset the current buffer length */   
            ctx->cur_buf_len = 0;
        }
    }

	return E_OK;
}

/************************************************************************************************************
 *  Cry_SM3_GetHash()
 ***********************************************************************************************************/
/*!
 * \brief         Finalizes the SM3 hash computation and retrieves the hash value.
 * \details       This function computes the final hash value by processing any remaining data in the buffer,
 *                appending padding, and incorporating the total message length. The final 256-bit hash is
 *                stored in the provided output buffer.
 *
 * \param[in,out] ctx        Pointer to the SM3 context structure.
 *                           - \a ctx->buf: Holds unprocessed data (if any) to be finalized.
 *                           - \a ctx->compressed_len: Tracks the total compressed data length in bits.
 *
 * \param[out]    output     Pointer to a buffer where the 32-byte (256-bit) hash result will be written.
 *
 * \return        E_OK       Operation completed successfully.
 * \return        E_NOT_OK   Invalid input parameters (e.g., NULL pointer or misaligned buffer).
 *
 * \pre           - \a ctx and \a output must be valid, non-NULL pointers.
 *                - \a ctx->buf and \a output must be properly 4-byte aligned.
 * \context       ANY
 * \reentrant     FALSE
 * \synchronous   TRUE
 * \config        -
 **********************************************************************************************************/
Std_ReturnType Cry_SM3_GetHash(sm3_ctx *ctx, uint8 *output)
{

    uint32 padn = 0;
    uint8 msglen[8] = {0};
    uint64 total_len = 0, high = 0, low = 0;

    // 1. Parameter validation
    if (ctx == NULL || output == NULL) 
    {
        return E_NOT_OK; // Invalid parameters
    }

    if (((uintptr_t)ctx->buf % 4 != 0) || ((uintptr_t)output % 4 != 0)) 
    {
        return E_NOT_OK; // Buffers must be 4-byte aligned
    }

    // 2. Calculate the total message length in bits
    total_len = ctx->compressed_len + ((uint64)ctx->cur_buf_len << 3);
    high = (total_len >> 32) & 0xFFFFFFFF;
    low = total_len & 0xFFFFFFFF;

    // Convert total message length to big-endian format
    PUT_ULONG_BE(high, msglen, 0);
    PUT_ULONG_BE(low, msglen, 4);

    // 3. Calculate padding length
    // Padding ensures the total length (including padding and length) is a multiple of 512 bits (64 bytes)
    padn = ((ctx->cur_buf_len + 1) <= 56) ? (56 - ctx->cur_buf_len) : (120 - ctx->cur_buf_len);

    // 4. Add padding and message length
    Cry_SM3_Update(ctx, (uint8 *)sm3_padding, padn); // Add padding
    Cry_SM3_Update(ctx, msglen, 8);                  // Add message length

    // 5. Write the final hash value (big-endian) to the output buffer
    for (int i = 0; i < 8; i++) 
    {
        PUT_ULONG_BE(ctx->state[i], output, i * 4);
    }

    // 6. Return success
    return E_OK;
}


/************************************************************************************************************
 *  Cry_SM3_MainFunction()
 ***********************************************************************************************************/
/*!
 * \brief         Performs the full SM3 hash computation on the input data.
 * \details       This function provides a simplified interface for computing the SM3 hash. It initializes 
 *                an SM3 context, updates it with the input data, and retrieves the final 256-bit hash value.
 *                It handles initialization, data processing, and finalization in a single call.
 *
 * \param[in]     input      Pointer to the input data to be hashed.
 *                           - This must be a valid, non-NULL pointer.
 *                           - Input data does not need to be aligned.
 *
 * \param[in]     iLen       The length of the input data (in bytes).
 *                           - This must be a positive integer and should not exceed the system's memory limits.
 *
 * \param[out]    output     Pointer to a buffer where the 32-byte (256-bit) hash result will be written.
 *                           - This must be a valid, non-NULL pointer, and the buffer must be at least 32 bytes.
 *
 * \return        Std_ReturnType
 *                - \c E_OK:     Operation completed successfully.
 *                - \c E_NOT_OK: One or more parameters are invalid (e.g., NULL pointer, zero length).
 *
 * \pre           - \a input and \a output must be valid, non-NULL pointers.
 *                - \a iLen must be greater than zero.
 * \context       ANY
 * \reentrant     FALSE
 * \synchronous   TRUE
 * \config        -
 **********************************************************************************************************/
Std_ReturnType Cry_SM3_MainFunction(const uint8 *input, uint32 iLen, uint8 *output)
{
    sm3_ctx mainctx; // SM3 context structure

    // 1. Parameter validation
    if (input == NULL || output == NULL) 
    {
        return E_NOT_OK; // Invalid parameters
    }
    // 2. Input length must be within valid range
    if (iLen == 0) 
    {
        return E_NOT_OK; // Input length must be within valid range
    }

    // 3. Initialize the SM3 context
    if (Cry_SM3_Init(&mainctx) != E_OK) 
    {
        return E_NOT_OK; // Initialization failed
    }

    // 4. Update the SM3 context with input data
    if (Cry_SM3_Update(&mainctx, input, iLen) != E_OK) 
    {
        return E_NOT_OK; // Update failed
    }

    // 5. Finalize the hash computation and retrieve the result
    if (Cry_SM3_GetHash(&mainctx, output) != E_OK) 
    {
        return E_NOT_OK; // Hash computation failed
    }

    return E_OK;
}
